Upgrade ONNX runtime, ONNX Gen AI runtime, add a new model to test changes by ammbra · Pull Request #1085 · openjdk/babylon

ammbra · 2026-07-03T08:38:19Z

I apologize for such a big PR, but most of its changes are result of regeneration for native bindings, operators, prototufs etc. So, this PR contains many changes that relate to upgrading the ONNX runtime library (1.26.0) and its native GenAI (0.14.0) counterpart:

The changes from oracle.code.onnx.foreign are generated via the cr-examples/onnx/opgen/setup.sh script and its helper symbols files to generate only the absolute necessary contents.
The operators and protobufs have been generated via their respective Java generators.
OnnxProtoBuilder needed to support also shapes because I kept getting this native error once rebuilding Babylon:

:bomb: Program crashed: Bad pointer dereference at 0x0000000000000008

Thread 37 "Java: oracle.code.onnx.llm.LlamaDemo.main()" crashed:

 0 0x000000012aa83f8c Generators::Model::IsPruned() const + 60 in libonnxruntime-genai.dylib
 1 0x000000012aa76cf4 Generators::Logits::Logits(Generators::State&) + 268 in libonnxruntime-genai.dylib
 2 0x000000012aa41848 Generators::DecoderOnly_State::DecoderOnly_State(Generators::DecoderOnly_Model const&, Generators::DeviceSpan<int>, Generators::GeneratorParams const&) + 88 in libonnxruntime-genai.dylib
 3 0x000000012aa4178c Generators::DecoderOnly_Model::CreateState(Generators::DeviceSpan<int>, Generators::GeneratorParams const&) const + 104 in libonnxruntime-genai.dylib
 4 0x000000012aa36e90 Generators::Generator::Generator(Generators::Model const&, Generators::GeneratorParams const&) + 432 in libonnxruntime-genai.dylib
 5 0x000000012aa36c90 Generators::CreateGenerator(Generators::Model const&, Generators::GeneratorParams const&) + 52 in libonnxruntime-genai.dylib
 6 0x000000012aade410 OgaCreateGenerator + 28 in libonnxruntime-genai.dylib
 7 0x000000013f6a2b00
 8 0x000000013f53cbc8
 9 0x000000013f53cbc8
10 0x000000013f53cbc8
11 0x000000013f53ca70
12 0x000000013f53ca70
13 0x000000013807e010
14 0x00007fffffffffff

I managed to identify that the exported model was the source for such an error (both local debug and by also comparing in Netron). To my understanding, Generators::Model::IsPruned() const method checks if the underlying model has been compressed using model pruning. Pruning seems to reduce model size and latency by zeroing out or permanently removing unnecessary weights and parameters, but I think in this case the generated model was aggressively pruned from the beginning of traversing the graph, thus resulting in the error above. The aggressive pruning was performed due the lack of shapes being there. This error occurred only for LLMs, the rest of the ML models being unaffected.

I also adjusted the license headers and cleaned a few leftovers in variable or patterns not really used by the library.
I updated the LlamaModel to use Reshape in order to fix tensor rank/shape visibility as the underlying tensor layout is already correct.
To test my changes beyond the existing demo, I created a new one. The AllMiniLML6V2EmbeddingModel is inspired by https://huggingface.co/onnx-community/all-MiniLM-L6-v2-ONNX and by running the OnnxLift over its default (model.onnx). As a result, I could obtain a quite good model to start. Then I cleaned it locally based on how the official config.json looked like. The BertTokenizer is just a helper for working with the tokenizer and its constants are from tokenizer_config.json and config.json , while the EmbeddingDemo is just a Java version of the code from the landing page of the model.

I confirm that I make this contribution in accordance with the OpenJDK Interim AI Policy.

Progress

Change must not contain extraneous whitespace

Reviewers

Adam Sotona (@asotona - Reviewer)

Reviewing

Using git

Checkout this PR locally:
$ git fetch https://git.openjdk.org/babylon.git pull/1085/head:pull/1085
$ git checkout pull/1085

Update a local copy of the PR:
$ git checkout pull/1085
$ git pull https://git.openjdk.org/babylon.git pull/1085/head

Using Skara CLI tools

Checkout this PR locally:
$ git pr checkout 1085

View PR using the GUI difftool:
$ git pr show -t 1085

Using diff file

Download this PR as a diff file:
https://git.openjdk.org/babylon/pull/1085.diff

Using Webrev

Link to Webrev Comment

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

bridgekeeper · 2026-07-03T08:40:31Z

👋 Welcome back amihalceanu! A progress list of the required criteria for merging this PR into code-reflection will be added to the body of your pull request. There are additional pull request commands available for use with this pull request.

openjdk · 2026-07-03T08:40:52Z

@ammbra This change now passes all automated pre-integration checks.

ℹ️ This project also has non-automated pre-integration requirements. Please see the file CONTRIBUTING.md for details.

After integration, the commit message for the final commit will be:

Upgrade ONNX runtime, ONNX Gen AI runtime, add a new model to test changes

Reviewed-by: asotona

You can use pull request commands such as /summary, /contributor and /issue to adjust it as needed.

At the time when this comment was updated there had been 15 new commits pushed to the code-reflection branch:

39496ab: Remove rogue character
e02dc39: Several boolean-returning code paths in BodyScanner lack target type
4192f1e: Fix conversion for compound assignment
... and 12 more: https://git.openjdk.org/babylon/compare/0cb9677f67964842dd046ad6bcf433ea2921667a...code-reflection

As there are no conflicts, your changes will automatically be rebased on top of these commits when integrating. If you prefer to avoid this automatic rebasing, please check the documentation for the /integrate command for further details.

As you do not have Committer status in this project an existing Committer must agree to sponsor your change. Possible candidates are the reviewers of this PR (@asotona) but any other Committer may sponsor as well.

➡️ To flag this PR as ready for integration with the above commit message, type /integrate in a new comment. (Afterwards, your sponsor types /sponsor in a new comment to perform the integration).

mlbridge · 2026-07-03T08:46:02Z

Webrevs

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

asotona

I’m still working through all the proposal, so please correct me if I misunderstand something.

There is clearly a need for named inputs and outputs in the graph. This is already handled inside OnnxTransformer, where ModuleAndInitializers holds the name map, allowing method parameter names and return record component names to be forwarded to the ONNX model.

The proposal seems to go in the opposite direction: to apply names parsed from the config file. I’m a bit concerned about this, because it moves us away from the idea that Java should remain the source of truth for the model.

There is also a clear need to provide shape information for some ONNX operations. This is new and would require passing additional information into the protobuf-building process. There are two parts to this issue:

identify the op input that requires an explicit shape
provide the actual shape value

If I understand the proposal correctly, the tensor would be identified by its name coming from the config, and the shape value would be represented as a list of objects, mixing literal values with references to the configuration file or even references to internal runtime values. My concern is that none of this information comes directly from the Java source.

If the shape definitions are primarily needed for method parameters, then the mapping key should probably be the tensor method parameter, or more generally some code-reflection value. If we had code reflection support for value attributes, or if we had support for parameter annotations, this could be handled more cleanly. Since we are not there yet, some workaround may be necessary. For example using parameter names as keys for the tensor shape resolver.

For the second part (the shape itself) I think it should still be Java-sourced. In practice, that would mean resolving it to an array of longs. Otherwise we risk introducing unclear references to “magic” constants coming from external configuration.

When I refer to shape inference, I mean that if shape information is required deeper in the model tree, it should ideally be inferred from the parameters and weights through the operations in the tree. Since that is a more complex task, we could temporarily rely on a Java variable-name-based mapping, even for values deeper in the model. Another option would be to identify the ops that require shape information and add synthetic shape attributes to them, in addition to their specs-generated attributes, so the shape constants are still passed explicitly from Java.

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

asotona

Nice work!
I’m glad there’s a clean workaround that replaces the shape hack.
Just a few minor comments below.

…ens. Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

ammbra · 2026-07-16T06:07:57Z

Thank you for your review, Adam. While addressing the issues noticed by you, I observed that the "fixed" model is slow. I was intrigued and looked into the model card example and noticed this for generating a response:

const output = await generator(messages, { max_new_tokens: 128 });

I further searched what that max_new_tokens stands for and found this https://huggingface.co/docs/transformers/main_classes/text_generation and this article https://medium.com/google-cloud/run-llms-anywhere-local-and-cpu-inference-with-onnx-runtime-genai-9bc34dbf0d7d.

Parameters: As you scale up max_new_tokens, the KV cache grows and the attention mechanism performs more computations. Keep an eye on memory usage and latency.

That quote matches the slow model at runtime behavior, once I introduced KV cache for Reshape usage.
So...my suspicion is that OnnxGenAI runtime defaults to max_length to calculate new token usage for a prompt, hence why KV cache usage grows.

For the moment, I added the max_new_tokens with value 128 in LlamaDemo and have a separate prompt method in OnnxGenAIRuntimeSession class to handle this way of working. Do you believe that max_new_tokens should be given as argument to LlamaDemo?

Otherwise, if current state of the implementation is ok, please help me integrate the changes.
/integrate

openjdk · 2026-07-16T06:10:24Z

@ammbra
Your change (at version 1b1c4a1) is now ready to be sponsored by a Committer.

asotona

I don't think creating a fresh generator in prompt is OK. It bypasses the session generator and drops the context between prompts.
maxNewTokens should be configured on the session generator

asotona · 2026-07-16T08:21:56Z

            return switch (this) {
-                case INT4, UINT4, FLOAT4E2M1 -> 4;
-                case UINT8, INT8, BOOL, FLOAT8E4M3FN, FLOAT8E4M3FNUZ, FLOAT8E5M2, FLOAT8E5M2FNUZ -> 8;
+                case INT2, INT4, UINT2, UINT4, FLOAT4E2M1 -> 4;


I think INT2 and UINT2 should have bit size 2, not 4.

You are right. I misread https://onnx.ai/onnx/technical/int2.html.

asotona · 2026-07-16T08:28:28Z

+    public void prompt(String prompt, int maxNewTokens, Consumer<String> outputConsumer) {
+        if (maxNewTokens < 0) {
+            throw new IllegalArgumentException("maxNewTokens must be above 0, found " + maxNewTokens);
+        }
+        LOG.log(System.Logger.Level.DEBUG, "Create sequences");
+        var inputTokens = call(OgaCreateSequences(ret));
+
+        MemorySegment localGeneratorParams = MemorySegment.NULL;
+        MemorySegment localGenerator = MemorySegment.NULL;
+        try {
+            call(OgaTokenizerEncode(tokenizer, arena.allocateFrom(prompt), inputTokens));
+            LOG.log(System.Logger.Level.DEBUG, "Tokenizer encode");
+            long promptTokens = OgaSequencesGetSequenceCount(inputTokens, 0L);
+            long maxLength =  Math.addExact(promptTokens, maxNewTokens);
+            localGeneratorParams = call(OgaCreateGeneratorParams(model, ret));
+            call(OgaGeneratorParamsSetSearchNumber(localGeneratorParams, arena.allocateFrom("max_length"), maxLength));
+            localGenerator = call(OgaCreateGenerator(model, localGeneratorParams, ret));
+
+            generate(inputTokens, localGenerator, outputConsumer);
        } finally {
+            if (!localGeneratorParams.equals(MemorySegment.NULL))
+                OgaDestroyGeneratorParams(localGeneratorParams);
+            if (!localGenerator.equals(MemorySegment.NULL))
+                OgaDestroyGenerator(localGenerator);


This looks like a workaround around the existing generator creation.

It is....I moved this logic in the main prompt generation, but this means that generator is no longer final and initialized with MemorySegment.NULL. generator depends on inputTokens (maxLength = promptTokens + maxNewTokens).

With a new generator the session loses context with each prompt. This makes the Llama session behave a bit like it has Alzheimer’s.

In that implementation...yes. So...I reverted the changes done there because what I actually did was to set the generator to use a max_length of prompt_length + 128. I did alter the genai_config.json max_length and is now set it for a lower value.

When max_length = context_length, the loop may continue much longer. As more tokens are generated (based on max_length = context_length), the KV cache grows and potentially that makes the model slower. I did try to set another provider (CoreML in my case) but the model is equally slow. Can you please take a look at the reverted changes and also the value set for max_length?

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

…other changes. Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

asotona · 2026-07-16T16:56:34Z

+    private final MemorySegment tokenizer;
+    private final MemorySegment tokenizerStream;
+    private final MemorySegment generatorParams;
+    private MemorySegment generator;


Could we make this field final again?

Yes, done. Thank you for catching this.

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

asotona

👍

ammbra · 2026-07-16T17:34:40Z

Thank you for all the time spent on this. Please
/integrate

openjdk · 2026-07-16T17:52:02Z

@ammbra
Your change (at version 494fcb9) is now ready to be sponsored by a Committer.

asotona · 2026-07-16T17:55:27Z

/sponsor

openjdk · 2026-07-16T17:56:14Z

Going to push as commit e109b0b.
Since your change was applied there have been 15 commits pushed to the code-reflection branch:

39496ab: Remove rogue character
e02dc39: Several boolean-returning code paths in BodyScanner lack target type
4192f1e: Fix conversion for compound assignment
... and 12 more: https://git.openjdk.org/babylon/compare/0cb9677f67964842dd046ad6bcf433ea2921667a...code-reflection

Your commit was automatically rebased without conflicts.

openjdk · 2026-07-16T17:56:26Z

@asotona @ammbra Pushed as commit e109b0b.

💡 You may see a message that your pull request was closed with unmerged commits. This can be safely ignored.

Upgrade ONNX runtime, gen ai and add bert model.

ec7ce5f

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

ammbra mentioned this pull request Jul 3, 2026

Upgrade ONNX runtime, ONNX Gen AI runtime, add a new model to test changes #1075

Closed

2 tasks

openjdk Bot added ready Pull request is ready to be integrated rfr Pull request is ready for review labels Jul 3, 2026

asotona reviewed Jul 3, 2026

View reviewed changes

ammbra added 4 commits July 4, 2026 19:11

Add a refactoring to how shapes are interpreted.

5ee2c38

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

Move constant to OnnxProtoBuilder

33a1227

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

Adjust package name upon refactoring

cc2308f

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

Replace Copyright with proper years.

48344f3

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

asotona reviewed Jul 7, 2026

View reviewed changes

ammbra and others added 2 commits July 14, 2026 11:47

Introduce another way to handle external shape metadata.

7f282b2

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

Revert work done on introducing shapes by adding Reshape to LlamaModel.

297d2cc

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

openjdk Bot removed ready Pull request is ready to be integrated rfr Pull request is ready for review labels Jul 14, 2026

ammbra added 3 commits July 14, 2026 15:53

Revert OnnxProtoBuilder state to basics.

d974000

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

Remove extra spaces.

e6ccad1

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

Fix white spaces in OnnxOps

76bf5bb

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

openjdk Bot added ready Pull request is ready to be integrated rfr Pull request is ready for review labels Jul 14, 2026

asotona approved these changes Jul 15, 2026

View reviewed changes

Comment thread cr-examples/onnx/opgen/src/main/java/oracle/code/onnx/Tensor.java Outdated

Comment thread cr-examples/onnx/src/main/java/oracle/code/onnx/genai/TensorDataStream.java Outdated

Comment thread cr-examples/onnx/src/main/java/oracle/code/onnx/ir/OnnxType.java Outdated

ammbra and others added 2 commits July 15, 2026 22:21

Address review comments and make the model faster by adding maxNewTok…

8ea7c36

…ens. Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

Fix rename typo for constant.

de0d005

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

openjdk Bot removed ready Pull request is ready to be integrated rfr Pull request is ready for review labels Jul 16, 2026

Fix trailing whitespace.

9baea9f

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

openjdk Bot added ready Pull request is ready to be integrated rfr Pull request is ready for review labels Jul 16, 2026

ammbra added 2 commits July 16, 2026 07:45

Update copyright as the file was modified in 2026 too

50df249

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

Refactor to avoid NPE for local generator params variable.

1b1c4a1

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

openjdk Bot added the sponsor Pull request is ready to be sponsored label Jul 16, 2026

asotona requested changes Jul 16, 2026

View reviewed changes

Use max_length in prompt response generation by default.

7b3ee76

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

openjdk Bot removed the sponsor Pull request is ready to be sponsored label Jul 16, 2026

Set max_length in genai_config.json to speed up the model and revert …

ed24160

…other changes. Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

asotona reviewed Jul 16, 2026

View reviewed changes

Make all MemorySegments final.

494fcb9

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

asotona approved these changes Jul 16, 2026

View reviewed changes

openjdk Bot added the sponsor Pull request is ready to be sponsored label Jul 16, 2026

openjdk Bot added the integrated Pull request has been integrated label Jul 16, 2026

openjdk Bot closed this Jul 16, 2026

openjdk Bot removed ready Pull request is ready to be integrated rfr Pull request is ready for review sponsor Pull request is ready to be sponsored labels Jul 16, 2026

Uh oh!

Conversation

ammbra commented Jul 3, 2026 • edited by openjdk Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Progress

Reviewers

Reviewing

Uh oh!

bridgekeeper Bot commented Jul 3, 2026

Uh oh!

openjdk Bot commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mlbridge Bot commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Webrevs

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

asotona left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

asotona left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ammbra commented Jul 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

openjdk Bot commented Jul 16, 2026

Uh oh!

asotona left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

asotona left a comment

Choose a reason for hiding this comment

Uh oh!

ammbra commented Jul 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

openjdk Bot commented Jul 16, 2026

Uh oh!

asotona commented Jul 16, 2026

Uh oh!

openjdk Bot commented Jul 16, 2026

Uh oh!

openjdk Bot commented Jul 16, 2026

Uh oh!

Reviewers

Assignees

ammbra commented Jul 3, 2026 •

edited by openjdk Bot

Loading

openjdk Bot commented Jul 3, 2026 •

edited

Loading

mlbridge Bot commented Jul 3, 2026 •

edited

Loading

asotona left a comment •

edited

Loading

ammbra commented Jul 16, 2026 •

edited

Loading

ammbra commented Jul 16, 2026 •

edited

Loading